# -*- encoding: utf-8 -*-
'''
@File    :   signal_processing.py
@Time    :   2022/3/22 14:20
@Author  :   ZhangChaoYang
@Desc    :   测试信号处理函数
'''

import os, sys

sys.path.insert(0, os.getcwd())
import unittest

import matplotlib.pyplot as plt
import numpy as np
from util.signal_processing import fft, point_pass, band_pass, multi_band_pass, cwt, sfft, ht, hht, hht_filter
from scipy.signal import chirp


class SPTest(unittest.TestCase):
    rps = 60  # 转速
    sample_rate = int(200 * rps)  # 采样频率
    duration = 0.05  # 采样持续几秒钟
    t = np.arange(0, duration, 1.0 / sample_rate)
    signal = None
    nonstationary_signal = None  # 非稳态信号

    @classmethod
    def setUpClass(cls) -> None:
        cls.f1 = 100
        cls.f2 = 300
        cls.f3 = 600
        # 模拟一个由多个正弦波合成的信号
        cls.part1 = np.cos(2 * np.pi * cls.f1 * cls.t)
        cls.part2 = 3 * np.cos(2 * np.pi * cls.f2 * cls.t)
        cls.part3 = 5 * np.cos(2 * np.pi * cls.f3 * cls.t)
        cls.signal = cls.part1 + cls.part2 + cls.part3
        # 模拟一个随时间变化的信号
        duration = 0.05  # 采样持续几秒钟
        cls.nonstationary_signal = np.piecewise(cls.t,
                                                [cls.t < duration, cls.t < 0.7 * duration, cls.t < 0.3 * duration],
                                                [lambda t: np.sin(2 * np.pi * cls.f1 * t),
                                                 lambda t: np.sin(2 * np.pi * cls.f2 * t),
                                                 lambda t: np.sin(2 * np.pi * cls.f3 * t)])

        plt.figure(1)
        plt.plot(cls.t, cls.signal, 'r')
        plt.suptitle("原始稳态波形")
        plt.xlabel('时间/秒')
        plt.ylabel('幅值')

        plt.figure(2)
        plt.plot(cls.t, cls.nonstationary_signal, 'b')
        plt.title("原始非稳态波形")
        plt.xlabel('时间/秒')
        plt.ylabel('幅值')

    @classmethod
    def setUp(self) -> None:
        plt.figure(3)

    @classmethod
    def tearDown(self) -> None:
        plt.show()

    def testFFT(self):
        xf, yf = fft(self.signal, self.sample_rate)
        plt.plot(xf, yf, 'b')
        plt.suptitle("傅利叶变换")
        plt.xlabel('频率/Hz')
        plt.ylabel('数量(归一化)')

    def testPointPass(self):
        freq = self.f1
        y = self.part1
        y_filtered = point_pass(self.signal, self.sample_rate, freq)
        plt.plot(self.t, y, 'r')
        plt.plot(self.t[:len(y_filtered)], y_filtered, 'b')
        plt.suptitle(f"frequency={freq}处点通滤波")
        plt.xlabel('时间/秒')
        plt.ylabel('幅值')

    def testBandPass(self):
        freq = self.f2
        y = self.part2
        plt.plot(self.t, y, 'r')
        y_filtered = band_pass(y, self.sample_rate, freq - 5, freq + 5)
        plt.plot(self.t[:len(y_filtered)], y_filtered, 'b')
        plt.suptitle(f"frequency={freq}附近带通滤波")
        plt.xlabel('时间/秒')
        plt.ylabel('幅值')

    def testMultiBandPass(self):
        freq1 = self.f1
        freq2 = self.f3
        y_filtereds = multi_band_pass(self.signal, self.sample_rate, [(freq1 - 5, freq1 + 5), (freq2 - 5, freq2 + 5)])
        for i, y_filtered in enumerate(y_filtereds):
            plt.plot(self.t[:len(y_filtered)], y_filtered, 'C' + str(i))
        plt.suptitle(f"frequency={freq1}和{freq2}附近带通滤波")
        plt.xlabel('时间/秒')
        plt.ylabel('幅值')

    def testCWT(self):
        xf, yf = cwt(self.signal, self.sample_rate)
        plt.plot(xf, yf, 'b')
        plt.suptitle("连续小波变换")
        plt.xlabel('频率/Hz')
        plt.ylabel('数量(归一化)')

    def testSFFT(self):
        time, f, magnitude = sfft(self.nonstationary_signal, self.sample_rate)
        plt.contourf(time, f, magnitude)
        plt.title('STFT Magnitude')
        plt.ylabel('频率/Hz')
        plt.xlabel('时间/秒')

    def testCWT2(self):
        f, magnitude = cwt(self.nonstationary_signal, self.sample_rate)
        plt.contourf(self.t, f, magnitude)
        plt.title('CWT Magnitude')
        plt.ylabel('频率/Hz')
        plt.xlabel('时间/秒')

    def testHT(self):
        rps = 60  # 转速
        sample_rate = int(200 * rps)  # 采样频率
        duration = 1  # 采样持续几秒钟
        t = np.arange(0, duration, 1.0 / sample_rate)
        origin_signal = chirp(t, 20.0, t[-1], 100.0)  # 原始信号
        y = origin_signal * (1.0 + 0.5 * np.sin(2.0 * np.pi * 3.0 * t))  # 对原始信号进行调制
        analytic_signal, instantaneous_frequency = ht(y, sample_rate)
        amplitude_envelope = np.abs(analytic_signal)  # 求绝对值得到幅值包络
        plt.plot(t, origin_signal, 'r')
        plt.plot(t, amplitude_envelope, 'b')
        plt.plot(t, y, 'g')
        plt.title('hilbert包络分析')
        plt.xlabel('时间/秒')
        plt.ylabel('幅值')

        plt.figure(3)
        plt.plot(t[:len(instantaneous_frequency)], instantaneous_frequency, 'b')
        plt.title('hilbert得到瞬时频率')
        plt.xlabel('时间/秒')
        plt.ylabel('频率/Hz')

    def testHHT(self):
        rps = 60  # 转速
        sample_rate = int(200 * rps)  # 采样频率
        duration = 1  # 采样持续几秒钟
        t = np.arange(0, duration, 1.0 / sample_rate)
        origin_signal = chirp(t, 20.0, t[-1], 100.0)  # 原始信号
        noise = np.piecewise(t, [t < duration, t < 0.7 * duration, t < 0.3 * duration],
                             [lambda t: np.sin(2 * np.pi * 5 * t), lambda t: np.sin(2 * np.pi * 10 * t),
                              lambda t: np.sin(2 * np.pi * 15 * t)])
        y = origin_signal * noise  # 对原始信号进行调制

        results = hht(y, sample_rate)
        n = len(results)
        _, axes = plt.subplots(n, 3, figsize=(20, 15))
        for i, result in enumerate(results):
            if i == 0:
                axes[i][0].set_title('原始信号')
            else:
                axes[i][0].set_title('IMF' + str(i))
            wave, tuple = result
            analytic_signal, instantaneous_frequency = tuple
            axes[i][0].plot(wave)
            axes[i][1].set_title('hilbert包络')
            axes[i][1].plot(np.abs(analytic_signal))
            axes[i][2].set_title('瞬时频率')
            axes[i][2].plot(instantaneous_frequency)

    def testHHTFilter(self):
        rps = 60  # 转速
        sample_rate = int(200 * rps)  # 采样频率
        duration = 1  # 采样持续几秒钟
        t = np.arange(0, duration, 1.0 / sample_rate)
        origin_signal = chirp(t, 20.0, t[-1], 100.0)  # 原始信号
        noise = np.piecewise(t, [t < duration, t < 0.7 * duration, t < 0.3 * duration],
                             [lambda t: np.sin(2 * np.pi * 5 * t), lambda t: np.sin(2 * np.pi * 10 * t),
                              lambda t: np.sin(2 * np.pi * 15 * t)])
        y = origin_signal * noise  # 对原始信号进行调制
        y_hat = hht_filter(y, sample_rate, 5, [0, 1, 2])
        plt.plot(t, y, 'r')
        plt.plot(t, y_hat, 'b')
        plt.title('EMD过滤')
        plt.xlabel('时间/秒')
        plt.ylabel('幅值')


if __name__ == "__main__":
    unittest.main()
    # suite = unittest.TestSuite()
    # suite.addTest(SPTest('testFFT'))
    # suite.addTest(SPTest('testPointPass'))
    # suite.addTest(SPTest('testBandPass'))
    # suite.addTest(SPTest('testMultiBandPass'))
    # suite.addTest(SPTest('testCWT'))
    # suite.addTest(SPTest('testSFFT'))
    # suite.addTest(SPTest('testCWT2'))
    # suite.addTest(SPTest('testHT'))
    # suite.addTest(SPTest('testHHT'))
    # suite.addTest(SPTest('testHHTFilter'))
    # runner = unittest.TextTestRunner()
    # runner.run(suite)
